# coding=utf-8

import logging
try:
    from cStringIO import StringIO
except ImportError:
    try:
        from StringIO import StringIO
    except ImportError:
        from io import StringIO

from paramiko.client import SSHClient, AutoAddPolicy
from paramiko.rsakey import RSAKey
import json

__all__ = ['Client']

logger = logging.getLogger('ssh.client')


class Client(object):
    def __init__(self, host, port=22):
        self.client = SSHClient()
        self.client.set_missing_host_key_policy(AutoAddPolicy())

        self.host = host
        self.port = port

    def connect(self, username, password):
        self.client.connect(self.host, self.port, username, password)

    def connect_private_key(self, username, private_key):
        buffer = StringIO(private_key)
        pkey = RSAKey.from_private_key(buffer)
        self.client.connect(self.host, self.port, username=username, pkey=pkey)

    def execute(self, command, data):
        stdin, stdout, stderr = self.client.exec_command(command, get_pty=True)
        data = json.dumps(data)

        is_stdin_close = False
        try:
            stdin.write(data)
            stdin.close()
            is_stdin_close = True
            result = stdout.read()
        except StandardError as e:
            logger.error('execute command error {0}'.format(e))
        finally:
            if not is_stdin_close:
                stdin.close()
            stdout.close()
            stderr.close()

        return result

    def close(self):
        self.client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

        if exc_type is None and exc_val is None and exc_tb is None:
            return False

        return True


if __name__ == '__main__':
    private_key = open('/Users/feipeng/.ssh/id_rsa_person').read()
    client = Client('45.63.123.226')
    client.connect_private_key('root', private_key=private_key)
    print(client.execute('/root/change_account', {'test': 'test'}))
